package com.katze.boot.applet.ssh;

import com.katze.boot.applet.ssh.service.SecureShellService;
import com.katze.common.ws.core.MessageChannel;
import com.katze.common.ws.core.MessageListener;
import com.katze.common.ws.core.UserChannelGroup;
import com.katze.common.ws.protocol.Message;
import com.katze.common.ws.protocol.MessageType;
import com.katze.ssh.cmd.ShellChannel;
import io.netty.channel.Channel;
import io.netty.channel.ChannelFutureListener;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.io.IOException;
import java.io.PipedOutputStream;
import java.io.Serializable;
import java.nio.charset.Charset;
import java.nio.charset.StandardCharsets;
import java.util.HashMap;
import java.util.Map;

@Component
@MessageListener(MessageType.SHELL)
public class ShellConsoleHandler implements MessageChannel {
    private static final Logger log = LoggerFactory.getLogger(ShellConsoleHandler.class);

    private final Map<Serializable, Map<String, ShellChannel>> maps = new HashMap<>();

    @Autowired
    private SecureShellService secureShellService;

    @Override
    public void onMessage(Channel channel, Message message) throws Exception {
        Map<String, ShellChannel> map = maps.computeIfAbsent(message.getSender(), (k) -> new HashMap<>(2));
        @SuppressWarnings("unchecked")
        Map<String, Object> data = (Map<String, Object>)message.getData();
        Object id = data.get("id");
        if (id == null){
            message.setCode(401);
            message.setData("未知会话，拒绝访问");
            channel.writeAndFlush(message).addListeners(ChannelFutureListener.CLOSE);
        } else {
            ShellChannel shellChannel = map.get((String) id);
            if (shellChannel == null) {
                shellChannel = secureShellService.createShell((String) id);
                shellChannel.setOut(msg -> channel.writeAndFlush(new Message((String) id, MessageType.SHELL, msg)));
                map.put((String) id, shellChannel);
            }
            switch ((String) data.get("cmd")) {
                case null:break;
                case "setup-charset":
                    shellChannel.changeCharset((String) data.get("charset"));
                    break;
                case "setup-pty":
                    @SuppressWarnings("unchecked")
                    Map<String,  Object> pty = (Map<String, Object>)data.get("pty");
                    shellChannel.changePtySize((Integer) pty.get("cols"), (Integer) pty.get("rows"));
                    break;
                case "exit":
                    shellChannel.close();
                    map.remove((String) id);
                    break;
                default:
                    shellChannel.send((String) data.get("cmd"));
            }
        }
    }

    @Override
    public void onClose(Channel channel) throws Exception {
        // websocket通道被关闭时，所有shell会话连接都将关闭
        Serializable user = channel.attr(UserChannelGroup.USER).get();
        Map<String, ShellChannel> map = maps.get(user);
        if (map != null) {
            for (ShellChannel value : map.values()) {
                value.close();
            }
            maps.remove( user);
            log.info("{}用户的Shell通道已关闭, 当前缓存数量: {}", user, maps.size());
        }
    }
}
